# coding: utf-8

import numpy as np
import wiringpi as wpi
from font import *


class OLed(object):
    def __init__(self, addr:int, screen_w:int=128, screen_h:int=64):
        assert screen_h % 8 == 0
        self.addr = addr
        self.screen_w = screen_w
        self.screen_h = screen_h
        self.device = wpi.wiringPiI2CSetup(self.addr)
        self.init()
        self.disbuf = np.zeros((screen_h//8, screen_w), dtype=np.uint8)


    def init(self):
        '''初始化'''
        self.writeCmd(0xAE)  # display off
        self.writeCmd(0x20)  # Set Memory Addressing Mode
        self.writeCmd(0x10)  # 00,Horizontal Addressing Mode;01,Vertical Addressing Mode;10,Page Addressing Mode (RESET);11,Invalid
        self.writeCmd(0xb0)  # Set Page Start Address for Page Addressing Mode,0-7
        self.writeCmd(0xc8)  # Set COM Output Scan Direction
        self.writeCmd(0x00)  # ---set low column address
        self.writeCmd(0x10)  # ---set high column address
        self.writeCmd(0x40)  # --set start line address
        self.writeCmd(0x81)  # --set contrast control register
        self.writeCmd(0xff)  # 亮度调节 0x00~0xff
        self.writeCmd(0xa1)  # --set segment re-map 0 to 127
        self.writeCmd(0xa6)  # --set normal display
        self.writeCmd(0xa8)  # --set multiplex ratio(1 to 64)
        self.writeCmd(0x3F)  #
        self.writeCmd(0xa4)  # 0xa4,Output follows RAM content;0xa5,Output ignores RAM content
        self.writeCmd(0xd3)  # -set display offset
        self.writeCmd(0x00)  # -not offset
        self.writeCmd(0xd5)  # --set display clock divide ratio/oscillator frequency
        self.writeCmd(0xf0)  # --set divide ratio
        self.writeCmd(0xd9)  # --set pre-charge period
        self.writeCmd(0x22)  #
        self.writeCmd(0xda)  # --set com pins hardware configuration
        self.writeCmd(0x12)
        self.writeCmd(0xdb)  # --set vcomh
        self.writeCmd(0x20)  # 0x20,0.77xVcc
        self.writeCmd(0x8d)  # --set DC-DC enable
        self.writeCmd(0x14)  #
        self.writeCmd(0xaf)  # --turn on oled panel


    def writeByte(self, addr:int, data:int):
        '''在指定位置写入数据'''
        # if isinstance(data, np.uint8):
        #     data = data.item()
        wpi.wiringPiI2CWriteReg8(self.device, addr, data)

    def writeCmd(self, cmd:int):
        '''写入命令'''
        self.writeByte(0x00, cmd)

    def writeData(self, data:int):
        '''写入数据'''
        self.writeByte(0x40, data)

    def setPos(self, x:int, y:int):
        '''设置位置'''
        self.writeCmd(0xB0 + y)     # page0-page7
        self.writeCmd(x&0x0F)       # low column start address
        self.writeCmd((x>>4)|0x10)  # high column start address

    def fill(self, c:int):
        '''填充屏幕'''
        buf_h, buf_w = self.disbuf.shape
        for y in range(buf_h):
            self.setPos(0, y)
            for x in range(buf_w):
                self.writeData(c)  # 写后列指针自动加一
                self.disbuf[y][x] = c

    def clear(self):
        '''清屏'''
        self.fill(0x00)

    def sleep(self):
        '''休眠'''
        self.writeCmd(0x8D)  # 设置电荷泵
        self.writeCmd(0x10)  # 关闭电荷泵
        self.writeCmd(0xAE)  # OLED休眠

    def wakeup(self):
        '''唤醒'''
        self.writeCmd(0x8D)  # 设置电荷泵
        self.writeCmd(0x14)  # 开启电荷泵
        self.writeCmd(0xAF)  # OLED唤醒

    def calcDisDiff(self, dis_data:list):
        '''计算要显示的数据与显存数据的差异
            dis_data数据格式: list[(addr, data),...]
        '''
        diff_data = []
        for (addr, data) in dis_data:
            x, y = addr
            if not data:
                continue
            for i, d in enumerate(data):
                if self.disbuf[y][x+i] != d:
                    diff_data.append([(x+i, y), d])
        return diff_data

    def batchDisData(self, dis_data:list, batch_interval:int=4):
        '''将临近地址的数据合并
            dis_data数据格式: list[(addr, data),...]
        '''
        diff_data_batch = []
        data_batch = None
        for (addr, data) in dis_data:
            if data_batch is None:
                data_batch = [addr, [data]]
            else:
                x, y = addr
                x_last, y_last = data_batch[0]
                x_last += len(data_batch[1]) - 1
                if y != y_last:
                    diff_data_batch.append(data_batch)
                    data_batch = [addr, [data]]
                elif x - x_last <= batch_interval:
                    interval = x - x_last
                    for i in range(1, interval):
                        data_batch[1].append(self.disbuf[y][x_last+i].item())
                    data_batch[1].append(data)
                else:
                    diff_data_batch.append(data_batch)
                    data_batch = [addr, [data]]
        if data_batch is not None:
            diff_data_batch.append(data_batch)
        return diff_data_batch

    def filterDisData(self, dis_data):
        '''过滤掉重复地址的数据'''
        dis_data_result = []
        addr_list = set()
        for i in range(len(dis_data)):
            i = len(dis_data) - i - 1
            addr = dis_data[i][0]
            if addr not in addr_list:
                addr_list.add(addr)
                dis_data_result.append(dis_data[i])
        return dis_data_result

    def showDisData(self, dis_data, need_fliter=True, need_sort=True):
        ''' 显示数据'''
        if need_fliter:
            dis_data = self.filterDisData(dis_data)
        if need_sort:
            dis_data.sort(key=lambda x: x[0][1]*self.screen_w+x[0][0])
        diff_data = self.calcDisDiff(dis_data)
        dis_data_batch = self.batchDisData(diff_data)
        for data_batch in dis_data_batch:
            x, y = data_batch[0]
            self.setPos(x, y)
            for i, d in enumerate(data_batch[1]):
                self.writeData(d)
                self.disbuf[y][x+i] = d

    def showString6x8(self, x:int, y:int, s:str, auto_new_line=True):
        '''显示6x8字符串'''
        dis_data = []
        for c in s:
            c = ord(c) - 32
            if x+6 > self.screen_w:
                if not auto_new_line:
                    break
                x = 0
                y += 1
                if y > 7:
                    y = 0
            data_line = [(x, y), []]
            for i in range(6):
                d = F6X8[c*6+i]
                data_line[1].append(d)
            dis_data.append(data_line)
            x += 6
        self.showDisData(dis_data)


    def showString8x16(self, x:int, y:int, s:str, auto_new_line=True):
        '''显示8x16字符串'''
        dis_data = []
        for c in s:
            c = ord(c) - 32
            if x+8 > self.screen_w:
                if not auto_new_line:
                    break
                x = 0
                y += 2
                if y > 7:
                    y = 0
            data_line = [(x, y), []]
            for i in range(8):
                d = F8X16[c*16+i]
                data_line[1].append(d)
            dis_data.append(data_line)
            data_line = [(x, y+1), []]
            for i in range(8):
                d = F8X16[c*16+i+8]
                data_line[1].append(d)
            dis_data.append(data_line)
            x += 8
        self.showDisData(dis_data)


    def showString(self, row:int, col:int, s:str, font:Font=Font6x8, auto_new_line=True):
        w = self.screen_w // font.width
        h = self.screen_h // font.height
        assert row < w and col < h
        x = row * font.width
        y = col * font.height // 8
        if font is Font6x8:
            self.showString6x8(x, y, s, auto_new_line)
        else:
            self.showString8x16(x, y, s, auto_new_line)


    def showImageData(self, data:np.array, x_offset:int=0, y_offset:int=0):
        '''显示图片数据'''
        h, w = data.shape
        dis_data = []
        for y in range(h):
            data_line = [(x_offset, y_offset+y), []]
            for x in range(w):
                data_line[1].append(data[y][x].item())
            dis_data.append(data_line)
        self.showDisData(dis_data, need_fliter=False, need_sort=False)


if __name__ == '__main__':
    import time
    oled = OLed(0x3C)
    oled.clear()

    time.sleep(1)
    for i in range(8):
        oled.showString(i, 1*i, f"{i+1}. Hello World! 0123456789",
                font=Font6x8, auto_new_line=False)
    time.sleep(1)
    for i in range(8):
        oled.showString(0, 1*i, f"{i+1}. Hello wORLD! 9876543210",
                font=Font6x8, auto_new_line=True)
    time.sleep(1)
    for i in range(4):
        oled.showString(0, 1*i, f"{i+1}. Hello wORLD! 9876543210",
                font=Font8x16, auto_new_line=False)
    time.sleep(1)
    oled.showString(0, 0, "[8x16] Hello World! 0123456789 0123456789",
            font=Font8x16, auto_new_line=False)
    time.sleep(1)
    oled.showString(0, 0, "[8x16] Hello World! 0123456789 0123456789",
            font=Font8x16)
    time.sleep(1)

    from img2bin import img2bindata
    data2 = img2bindata('test.jpg', resize=(128, 64))
    oled.showImageData(data2)
